import os

import torch.nn as nn
from torch import optim
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import random
from torchvision import datasets, transforms

import numpy as np
import torch
import os
import matplotlib.pyplot as plt
from LeNet5卷神经网络 import LeNet5

#设置随机数种子，保证结果的可复现
def setup_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)

#设置随机数种子
setup_seed(26)

#加载数据,ToTensor()将数据转换为tensor
train_data = datasets.MNIST(root='./data', train=True,download=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data',train=False,download=True, transform=transforms.ToTensor())

#数据切分
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)

model = LeNet5().to('cpu')

#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

total = 0
corrt = 0
model.load_state_dict(torch.load('./model.pth'))
#临时禁用自动梯度计算，验证时只关心输出
#节省内存、加速计算、防止梯度累加
with torch.no_grad():
    model.eval()
    for images, labels in test_loader:
        images = images.to('cpu')
        labels = labels.to('cpu')
        outputs = model(images)
        # print(outputs)#[64,10]/[batch_size,labels]/[一次迭代中模型处理的样本数量, 10 个类别]

        _,predicted = torch.max(outputs, 1)

        #求准确率
        #求样本数量
        total += labels.size(0)
        #预测正确的数量
        corrt += (predicted == labels).sum().item()

print(f'准确率:{100 * corrt / total}%')







